from flask import Blueprint, render_template, redirect, url_for, flash, request, send_file, current_app
from flask_login import login_required, current_user
from ..models.pathogen import Pathogen
from ..forms import PathogenForm, PathogenSearchForm
from .. import db
from datetime import datetime, timedelta
from sqlalchemy import or_
from collections import defaultdict
import json
from Bio.Seq import Seq
from Bio import pairwise2
import numpy as np
from collections import Counter
from Bio import Align  # 使用新的对齐器
from docx import Document
from docx.shared import Inches
import matplotlib.pyplot as plt
import io
import os

# 创建蓝图
pathogen_bp = Blueprint('pathogen', __name__)

@pathogen_bp.route('/list')
@login_required
def list():
    form = PathogenSearchForm(request.args)
    
    # 构建查询
    query = Pathogen.query
    
    # 关键词搜索
    if request.args.get('keyword'):
        keyword = f"%{request.args.get('keyword')}%"
        query = query.filter(or_(
            Pathogen.name.like(keyword),
            Pathogen.isolation_location.like(keyword),
            Pathogen.gene_sequence.like(keyword)
        ))
    
    # 类型过滤
    if request.args.get('type'):
        query = query.filter(Pathogen.type == request.args.get('type'))
    
    # 日期范围过滤
    if request.args.get('date_from'):
        date_from = datetime.strptime(request.args.get('date_from'), '%Y-%m-%d')
        query = query.filter(Pathogen.isolation_date >= date_from)
    if request.args.get('date_to'):
        date_to = datetime.strptime(request.args.get('date_to'), '%Y-%m-%d')
        query = query.filter(Pathogen.isolation_date <= date_to)
    
    # 排序
    sort = request.args.get('sort', 'created_at')
    order = request.args.get('order', 'desc')
    
    if order == 'desc':
        query = query.order_by(getattr(Pathogen, sort).desc())
    else:
        query = query.order_by(getattr(Pathogen, sort).asc())
    
    # 分页
    page = request.args.get('page', 1, type=int)
    pathogens = query.paginate(page=page, per_page=10, error_out=False)
    
    return render_template('pathogen/list.html', 
                         pathogens=pathogens,
                         form=form,
                         sort=sort,
                         order=order)

@pathogen_bp.route('/new', methods=['GET', 'POST'])
@login_required
def create():
    form = PathogenForm()
    if form.validate_on_submit():
        pathogen = Pathogen(
            name=form.name.data,
            type=form.type.data,
            isolation_date=form.isolation_date.data,
            isolation_location=form.isolation_location.data,
            latitude=form.latitude.data,
            longitude=form.longitude.data,
            gene_sequence=form.gene_sequence.data,
            created_by=current_user.id
        )
        db.session.add(pathogen)
        db.session.commit()
        flash('病原体数据添加成功', 'success')
        return redirect(url_for('pathogen.list'))
    return render_template('pathogen/form.html', form=form, title='新增病原体')

@pathogen_bp.route('/<int:id>/edit', methods=['GET', 'POST'])
@login_required
def edit(id):
    pathogen = Pathogen.query.get_or_404(id)
    form = PathogenForm(obj=pathogen)
    if form.validate_on_submit():
        pathogen.name = form.name.data
        pathogen.type = form.type.data
        pathogen.isolation_date = form.isolation_date.data
        pathogen.isolation_location = form.isolation_location.data
        pathogen.latitude = form.latitude.data
        pathogen.longitude = form.longitude.data
        pathogen.gene_sequence = form.gene_sequence.data
        db.session.commit()
        flash('病原体数据更新成功', 'success')
        return redirect(url_for('pathogen.list'))
    return render_template('pathogen/form.html', form=form, title='编辑病原体')

@pathogen_bp.route('/<int:id>/delete', methods=['POST'])
@login_required
def delete(id):
    pathogen = Pathogen.query.get_or_404(id)
    db.session.delete(pathogen)
    db.session.commit()
    flash('病原体数据已删除', 'success')
    return redirect(url_for('pathogen.list'))

@pathogen_bp.route('/analysis')
@login_required
def analysis():
    # 获取所有病原体数据
    pathogens = Pathogen.query.all()
    
    # 按类型统计数量
    type_stats = db.session.query(
        Pathogen.type, 
        db.func.count(Pathogen.id)
    ).group_by(Pathogen.type).all()
    
    # 按月份统计数量
    month_stats = db.session.query(
        db.func.date_format(Pathogen.isolation_date, '%Y-%m'),
        db.func.count(Pathogen.id)
    ).group_by(db.func.date_format(Pathogen.isolation_date, '%Y-%m')
    ).order_by(db.func.date_format(Pathogen.isolation_date, '%Y-%m')).all()
    
    # 按地区统计
    location_stats = defaultdict(lambda: {'count': 0, 'types': defaultdict(int)})
    for p in pathogens:
        location_stats[p.isolation_location]['count'] += 1
        location_stats[p.isolation_location]['types'][p.type] += 1
    
    return render_template('pathogen/analysis.html',
                         type_stats=type_stats,
                         month_stats=month_stats,
                         location_stats=dict(location_stats))

@pathogen_bp.route('/spread')
@login_required
def spread():
    # 获取时间范围
    end_date = datetime.now()
    start_date = end_date - timedelta(days=365)  # 默认分析最近一年
    
    # 获取所有病原体数据
    pathogens = Pathogen.query.filter(
        Pathogen.isolation_date.between(start_date, end_date)
    ).order_by(Pathogen.isolation_date).all()
    
    # 按地区统计
    region_stats = defaultdict(lambda: {'count': 0, 'types': defaultdict(int)})
    for p in pathogens:
        region_stats[p.isolation_location]['count'] += 1
        region_stats[p.isolation_location]['types'][p.type] += 1
    
    # 时空传播数据
    spread_data = []
    for p in pathogens:
        if p.latitude and p.longitude:
            spread_data.append({
                'name': p.name,
                'type': p.type,
                'location': p.isolation_location,
                'date': p.isolation_date.strftime('%Y-%m-%d'),
                'coordinates': [p.latitude, p.longitude]
            })
    
    # 关联分析
    type_correlations = {}
    for type1 in ['virus', 'bacteria', 'parasite', 'fungus']:
        type_correlations[type1] = {}
        for type2 in ['virus', 'bacteria', 'parasite', 'fungus']:
            if type1 != type2:
                # 计算两种类型在同一地区出现的次数
                coexist_count = db.session.query(Pathogen).filter(
                    Pathogen.isolation_location.in_(
                        db.session.query(Pathogen.isolation_location).filter(
                            Pathogen.type == type1
                        )
                    ),
                    Pathogen.type == type2
                ).count()
                type_correlations[type1][type2] = coexist_count
    
    return render_template('pathogen/spread.html',
                         pathogens=pathogens,
                         region_stats=dict(region_stats),
                         spread_data=json.dumps(spread_data),
                         type_correlations=type_correlations)

@pathogen_bp.route('/mutation')
@login_required
def mutation():
    # 获取所有病原体数据
    pathogens = Pathogen.query.all()
    
    # 按类型分组基因序列 - 修改这里
    sequences_by_type = {}  # 使用普通字典
    
    # 初始化每种类型的列表
    for type_name in ['virus', 'bacteria', 'parasite', 'fungus']:
        sequences_by_type[type_name] = []
    
    # 添加序列数据
    for p in pathogens:
        if p.gene_sequence and p.type in sequences_by_type:  # 确保有基因序列且类型有效
            sequences_by_type[p.type].append({
                'name': p.name,
                'sequence': p.gene_sequence,
                'date': p.isolation_date
            })
    
    # 计算每种类型的序列相似度
    similarity_stats = {}
    for type_name, sequences in sequences_by_type.items():
        if len(sequences) > 1:
            similarities = []
            for i in range(len(sequences)):
                for j in range(i + 1, len(sequences)):
                    seq1 = Seq(sequences[i]['sequence'])
                    seq2 = Seq(sequences[j]['sequence'])
                    # 计算序列相似度
                    aligner = Align.PairwiseAligner()
                    aligner.mode = 'global'
                    aligner.match_score = 1.0
                    aligner.mismatch_score = 0.0
                    aligner.gap_score = -0.5
                    alignments = aligner.align(seq1, seq2)
                    if alignments:
                        similarity = alignments[0].score / max(len(seq1), len(seq2))
                        similarities.append(similarity)
            
            if similarities:
                similarity_stats[type_name] = {
                    'avg': float(np.mean(similarities)),
                    'min': float(np.min(similarities)),
                    'max': float(np.max(similarities))
                }
    
    # 计算碱基组成
    base_compositions = {}
    for type_name, sequences in sequences_by_type.items():
        all_bases = ''.join(s['sequence'] for s in sequences)
        base_count = Counter(all_bases)
        total = sum(base_count.values())
        if total > 0:
            base_compositions[type_name] = {
                base: float(count/total) for base, count in base_count.items()
            }
    
    # 计算时间序列变异
    time_variations = {}
    for type_name, sequences in sequences_by_type.items():
        if len(sequences) > 1:
            # 按时间排序
            sorted_seqs = sorted(sequences, key=lambda x: x['date'])
            variations = []
            reference = sorted_seqs[0]['sequence']  # 使用最早的序列作为参考
            for seq in sorted_seqs[1:]:
                # 计算与参考序列的差异
                diff_count = sum(1 for a, b in zip(reference, seq['sequence']) if a != b)
                variations.append({
                    'date': seq['date'].strftime('%Y-%m'),
                    'diff_rate': float(diff_count / len(reference)) if len(reference) > 0 else 0.0
                })
            time_variations[type_name] = variations
    
    return render_template('pathogen/mutation.html',
                         similarity_stats=similarity_stats,
                         base_compositions=base_compositions,
                         time_variations=time_variations)

@pathogen_bp.route('/report')
@login_required
def generate_report():
    # 获取所有病原体数据
    pathogens = Pathogen.query.all()
    
    # 创建Word文档
    doc = Document()
    
    # 添加标题
    doc.add_heading('滩羊主要地方流行性疫病病原区系分析报告', 0)
    doc.add_paragraph(f'生成时间：{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}')
    
    # 1. 基本统计信息
    doc.add_heading('1. 基本统计信息', level=1)
    # 按类型统计
    type_stats = {}
    for p in pathogens:
        type_stats[p.type] = type_stats.get(p.type, 0) + 1
    
    table = doc.add_table(rows=1, cols=2)
    table.style = 'Table Grid'
    table.rows[0].cells[0].text = '病原体类型'
    table.rows[0].cells[1].text = '数量'
    for type_name, count in type_stats.items():
        row = table.add_row()
        row.cells[0].text = type_name
        row.cells[1].text = str(count)
    
    # 2. 地理分布分析
    doc.add_heading('2. 地理分布分析', level=1)
    location_stats = {}
    for p in pathogens:
        location_stats[p.isolation_location] = location_stats.get(p.isolation_location, 0) + 1
    
    table = doc.add_table(rows=1, cols=2)
    table.style = 'Table Grid'
    table.rows[0].cells[0].text = '采样地点'
    table.rows[0].cells[1].text = '数量'
    for location, count in location_stats.items():
        row = table.add_row()
        row.cells[0].text = location
        row.cells[1].text = str(count)
    
    # 3. 变异分析
    doc.add_heading('3. 变异分析', level=1)
    
    # 按类型分组基因序列
    sequences_by_type = {}
    for type_name in ['virus', 'bacteria', 'parasite', 'fungus']:
        sequences_by_type[type_name] = []
    
    for p in pathogens:
        if p.gene_sequence and p.type in sequences_by_type:
            sequences_by_type[p.type].append({
                'name': p.name,
                'sequence': p.gene_sequence,
                'date': p.isolation_date
            })
    
    # 计算变异统计
    for type_name, sequences in sequences_by_type.items():
        if len(sequences) > 1:
            doc.add_heading(f'{type_name}变异分析', level=2)
            p = doc.add_paragraph()
            p.add_run(f'样本数量：{len(sequences)}个\n')
            
            # 计算序列相似度
            similarities = []
            for i in range(len(sequences)):
                for j in range(i + 1, len(sequences)):
                    seq1 = Seq(sequences[i]['sequence'])
                    seq2 = Seq(sequences[j]['sequence'])
                    aligner = Align.PairwiseAligner()
                    aligner.mode = 'global'
                    aligner.match_score = 1.0
                    aligner.mismatch_score = 0.0
                    aligner.gap_score = -0.5
                    alignments = aligner.align(seq1, seq2)
                    if alignments:
                        similarity = alignments[0].score / max(len(seq1), len(seq2))
                        similarities.append(similarity)
            
            if similarities:
                p.add_run(f'平均相似度：{(np.mean(similarities) * 100):.1f}%\n')
                p.add_run(f'最大相似度：{(np.max(similarities) * 100):.1f}%\n')
                p.add_run(f'最小相似度：{(np.min(similarities) * 100):.1f}%\n')
    
    # 4. 结论与建议
    doc.add_heading('4. 结论与建议', level=1)
    doc.add_paragraph('根据以上分析结果，我们得出以下结论：')
    conclusions = doc.add_paragraph(style='List Bullet')
    conclusions.add_run('1. 病原体分布特点：\n')
    for type_name, count in type_stats.items():
        if count > 0:
            conclusions.add_run(f'   - {type_name}：发现{count}个样本\n')
    
    conclusions.add_run('\n2. 地理分布特点：\n')
    for location, count in location_stats.items():
        conclusions.add_run(f'   - {location}：发现{count}个样本\n')
    
    # Windows 风格的路径处理
    report_path = os.path.join(current_app.root_path, 'static', 'reports')
    report_path = os.path.normpath(report_path)  # 标准化路径
    if not os.path.exists(report_path):
        os.makedirs(report_path)
    
    filename = f'pathogen_report_{datetime.now().strftime("%Y%m%d_%H%M%S")}.docx'
    full_path = os.path.join(report_path, filename)
    full_path = os.path.normpath(full_path)  # 标准化路径
    
    # 确保目录存在
    os.makedirs(os.path.dirname(full_path), exist_ok=True)
    
    # 保存文档
    doc.save(full_path)
    
    # 返回文件下载链接
    return send_file(
        full_path,
        as_attachment=True,
        download_name=filename,
        mimetype='application/vnd.openxmlformats-officedocument.wordprocessingml.document'
    ) 